package com.sig.agent;

import com.sun.tools.attach.VirtualMachine;

import java.io.File;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * 热更任务
 *
 * @author sig
 * @since 2019/3/6
 */
public class HotUpdateTask extends Thread {
    private final String agentJar;
    private final String reloadPath;
    private long period;
    private long keep;
    private final String pid;
    private final VirtualMachine vm;

    private volatile boolean running = true;
    private final Map<String, ClassFileUpdate> classFileUpdates = new HashMap<>();

    public HotUpdateTask(String agentJar, String reloadPath, long period, long keep, boolean needDelete) throws Exception {
        this.agentJar = agentJar;
        this.reloadPath = reloadPath;
        this.period = period;
        this.keep = keep;
        this.pid = ManagementFactory.getRuntimeMXBean().getName().split("@")[0];
        this.vm = VirtualMachine.attach(pid);
        if (needDelete && keep > 0L)
            this.deleteReloadClass();
    }

    public long getPeriod() {
        return period;
    }

    public void setPeriod(long period) {
        this.period = period;
    }

    public long getKeep() {
        return keep;
    }

    public void setKeep(long keep) {
        this.keep = keep;
    }

    public void shutdown() {
        running = false;
        System.out.println("Hot update shutdown:" + this);
    }

    public void destroy() {
        shutdown();
        classFileUpdates.clear();
        if (vm != null) {
            try {
                vm.detach();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        System.out.println("Hot update destroy:" + this);
    }

    @Override
    public void start() {
        super.start();
        System.out.println("Hot update start:" + this);
    }

    @Override
    public void run() {
        while (running) {
            try {
                if (period > 0L)
                    TimeUnit.MILLISECONDS.sleep(period);
                checkClasses();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    private void checkClasses() {
        File file = classesPath();
        File[] files = file.listFiles();
        if (files == null)
            return;
        boolean success = false;
        for (File tmpFile : files) {
            String fileName = tmpFile.getName();
            if (!fileName.toLowerCase().endsWith(".class"))
                continue;
            ClassFileUpdate classFileUpdate = classFileUpdates.get(fileName);
            if (classFileUpdate == null) {
                classFileUpdate = new ClassFileUpdate(fileName);
                classFileUpdates.put(fileName, classFileUpdate);
            }
            long now = System.currentTimeMillis();
            if (file.lastModified() == classFileUpdate.lastModifyTime()) {
                if (keep > 0L && now - classFileUpdate.lastReloadTime() >= keep) {
                    file.delete();
                }
            } else {
                String classFilePath = file.getPath();
                if (updateClass(classFilePath)) {
                    System.out.println("Hot update:" + classFilePath + " success, cost(ms):" + (System.currentTimeMillis() - now));
                    classFileUpdate.add2UpdateHistory(file.lastModified());
                    success = true;
                } else {
                    System.out.println("Hot update:" + classFilePath + "failed");
                }
            }
        }
        if (success)
            System.out.println("Hot updates:" + classFileUpdates);
    }

    private boolean updateClass(String classFilePath) {
        try {
            vm.loadAgent(agentJar, classFilePath);
            return true;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return false;
    }

    private void deleteReloadClass() {
        File file = classesPath();
        File[] files = file.listFiles();
        if (files != null) {
            for (File tmpFile : files)
                tmpFile.delete();
        }
    }

    private File classesPath() {
        return new File(reloadPath);
    }

    protected boolean running() {
        return running;
    }

    @Override
    public String toString() {
        StringBuilder builder = new StringBuilder();
        builder.append("HotUpdateTask [")
                .append("pid=").append(pid)
                .append(", period=").append(period)
                .append(", keep=").append(keep)
                .append(", agentJar=").append(agentJar)
                .append(", reloadPath=").append(reloadPath)
                .append("]");
        return builder.toString();
    }
}
